#include <gtest/gtest.h>
#include "utils_task.h"
#include <chrono>
#include <atomic>
#include <thread>
#include <vector>

using namespace El::Utils;

class CommonTaskTest : public ::testing::Test {
protected:
    void SetUp() override {
        // Reset singleton for each test
        TaskScheduler::GetInstance().SetThreadPoolSize(2);
    }

    void TearDown() override {
        // Clean up after each test
    }
};

// Test singleton pattern
TEST_F(CommonTaskTest, SingletonTest) {
    TaskScheduler& instance1 = TaskScheduler::GetInstance();
    TaskScheduler& instance2 = TaskScheduler::GetInstance();
    EXPECT_EQ(&instance1, &instance2);
}

// Test thread pool size setting
TEST_F(CommonTaskTest, ThreadPoolSizeTest) {
    TaskScheduler& task = TaskScheduler::GetInstance();
    task.SetThreadPoolSize(4);
    // Add a task to verify the thread pool is working
    std::atomic<bool> taskExecuted{false};
    auto future = task.PostTask("test", [&taskExecuted]() {
        taskExecuted = true;
    });
    future.wait();
    EXPECT_TRUE(taskExecuted);
}

// Test basic task execution
TEST_F(CommonTaskTest, BasicTaskExecutionTest) {
    TaskScheduler& task = TaskScheduler::GetInstance();
    std::atomic<int> counter{0};
    
    auto future = task.PostTask("test", [&counter]() {
        counter++;
    });
    
    future.wait();
    EXPECT_EQ(counter, 1);
}

// Test multiple tasks with same name
TEST_F(CommonTaskTest, MultipleTasksSameNameTest) {
    TaskScheduler& task = TaskScheduler::GetInstance();
    std::atomic<int> counter{0};
    
    std::vector<std::future<void>> futures;
    for (int i = 0; i < 5; ++i) {
        futures.push_back(task.PostTask("same_name", [&counter]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
            counter++;
        }));
    }
    
    // Wait for all tasks to complete
    for (auto& future : futures) {
        future.wait();
    }
    
    EXPECT_EQ(counter, 5);
}

// Test multiple tasks with different names
TEST_F(CommonTaskTest, MultipleTasksDifferentNamesTest) {
    TaskScheduler& task = TaskScheduler::GetInstance();
    std::atomic<int> counter{0};
    
    std::vector<std::future<void>> futures;
    for (int i = 0; i < 5; ++i) {
        futures.push_back(task.PostTask("task_" + std::to_string(i), [&counter]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
            counter++;
        }));
    }
    
    // Wait for all tasks to complete
    for (auto& future : futures) {
        future.wait();
    }
    
    EXPECT_EQ(counter, 5);
}

// Test task execution order within same task name
TEST_F(CommonTaskTest, TaskOrderTest) {
    TaskScheduler& task = TaskScheduler::GetInstance();
    std::vector<int> executionOrder;
    std::mutex orderMutex;
    
    std::vector<std::future<void>> futures;
    for (int i = 0; i < 3; ++i) {
        futures.push_back(task.PostTask("ordered_task", [i, &executionOrder, &orderMutex]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
            std::lock_guard<std::mutex> lock(orderMutex);
            executionOrder.push_back(i);
        }));
    }
    
    // Wait for all tasks to complete
    for (auto& future : futures) {
        future.wait();
    }
    
    // Check if tasks were executed in order
    EXPECT_EQ(executionOrder.size(), (size_t)3);
    for (size_t i = 0; i < executionOrder.size(); ++i) {
        EXPECT_EQ(executionOrder[i], i);
    }
}

// Test concurrent task execution with different names
TEST_F(CommonTaskTest, ConcurrentExecutionTest) {
    TaskScheduler& task = TaskScheduler::GetInstance();
    task.SetThreadPoolSize(4);  // Set more threads for better concurrency
    
    std::atomic<int> counter{0};
    std::chrono::steady_clock::time_point start = std::chrono::steady_clock::now();
    
    std::vector<std::future<void>> futures;
    for (int i = 0; i < 4; ++i) {
        futures.push_back(task.PostTask("task_" + std::to_string(i), [&counter]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            counter++;
        }));
    }
    
    // Wait for all tasks to complete
    for (auto& future : futures) {
        future.wait();
    }
    
    std::chrono::steady_clock::time_point end = std::chrono::steady_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    
    EXPECT_EQ(counter, 4);
    // If tasks run concurrently, total time should be less than 400ms (4 * 100ms)
    EXPECT_LT(duration, 400);
}

// Test sequential execution of same-name tasks
TEST_F(CommonTaskTest, SequentialExecutionSameNameTest) {
    TaskScheduler& task = TaskScheduler::GetInstance();
    std::vector<int> executionOrder;
    std::mutex orderMutex;
    
    // Create 5 tasks with the same name
    std::vector<std::future<void>> futures;
    for (int i = 0; i < 5; ++i) {
        futures.push_back(task.PostTask("sequential_task", [i, &executionOrder, &orderMutex]() {
            // Simulate some work
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
            std::lock_guard<std::mutex> lock(orderMutex);
            executionOrder.push_back(i);
        }));
    }
    
    // Wait for all tasks to complete
    for (auto& future : futures) {
        future.wait();
    }
    
    // Verify tasks were executed in order
    ASSERT_EQ(executionOrder.size(), 5);
    for (size_t i = 0; i < executionOrder.size(); ++i) {
        EXPECT_EQ(executionOrder[i], i) << "TaskScheduler " << i << " executed out of order";
    }
}

// Test parallel execution of different-name tasks
TEST_F(CommonTaskTest, ParallelExecutionDifferentNamesTest) {
    TaskScheduler& task = TaskScheduler::GetInstance();
    task.SetThreadPoolSize(4);  // Set more threads for better parallelism
    
    std::vector<std::chrono::steady_clock::time_point> startTimes;
    std::vector<std::chrono::steady_clock::time_point> endTimes;
    std::mutex timesMutex;
    
    // Create tasks with different names
    std::vector<std::future<void>> futures;
    for (int i = 0; i < 4; ++i) {
        futures.push_back(task.PostTask("task_" + std::to_string(i), [i, &startTimes, &endTimes, &timesMutex]() {
            {
                std::lock_guard<std::mutex> lock(timesMutex);
                startTimes.push_back(std::chrono::steady_clock::now());
            }
            // Each task takes 100ms to complete
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            {
                std::lock_guard<std::mutex> lock(timesMutex);
                endTimes.push_back(std::chrono::steady_clock::now());
            }
        }));
    }
    
    // Wait for all tasks to complete
    for (auto& future : futures) {
        future.wait();
    }
    
    // Verify parallel execution
    ASSERT_EQ(startTimes.size(), 4);
    ASSERT_EQ(endTimes.size(), 4);
    
    // Calculate the total execution time
    auto firstStart = *std::min_element(startTimes.begin(), startTimes.end());
    auto lastEnd = *std::max_element(endTimes.begin(), endTimes.end());
    auto totalDuration = std::chrono::duration_cast<std::chrono::milliseconds>(lastEnd - firstStart).count();
    
    // If tasks run in parallel, total time should be around 100ms
    // If tasks run sequentially, it would take about 400ms
    EXPECT_LT(totalDuration, 200) << "Tasks appear to be running sequentially instead of in parallel";
}

// Test mixed sequential and parallel execution
TEST_F(CommonTaskTest, MixedSequentialAndParallelTest) {
    TaskScheduler& task = TaskScheduler::GetInstance();
    task.SetThreadPoolSize(4);
    
    std::vector<std::pair<std::string, int>> executionOrder;
    std::mutex orderMutex;
    auto recordExecution = [&executionOrder, &orderMutex](const std::string& taskName, int index) {
        std::lock_guard<std::mutex> lock(orderMutex);
        executionOrder.push_back({taskName, index});
    };
    
    // Create two sets of tasks: 3 tasks with name "A" and 3 tasks with name "B"
    std::vector<std::future<void>> futures;
    
    // Push tasks alternately
    for (int i = 0; i < 3; ++i) {
        // TaskScheduler A
        futures.push_back(task.PostTask("A", [i, &recordExecution]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
            recordExecution("A", i);
        }));
        
        // TaskScheduler B
        futures.push_back(task.PostTask("B", [i, &recordExecution]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
            recordExecution("B", i);
        }));
    }
    
    // Wait for all tasks to complete
    for (auto& future : futures) {
        future.wait();
    }
    
    // Extract sequences for each task name
    std::vector<int> sequenceA, sequenceB;
    for (const auto& execution : executionOrder) {
        if (execution.first == "A") {
            sequenceA.push_back(execution.second);
        } else {
            sequenceB.push_back(execution.second);
        }
    }
    
    // Verify that each sequence is in order
    ASSERT_EQ(sequenceA.size(), 3);
    ASSERT_EQ(sequenceB.size(), 3);
    
    // Check if tasks with the same name were executed in order
    for (size_t i = 0; i < sequenceA.size(); ++i) {
        EXPECT_EQ(sequenceA[i], i) << "TaskScheduler A" << i << " executed out of order";
        EXPECT_EQ(sequenceB[i], i) << "TaskScheduler B" << i << " executed out of order";
    }
    
    // Verify that total execution time indicates parallel execution between A and B
    EXPECT_EQ(executionOrder.size(), 6);
}